# -*- coding: utf-8 -*-

from PyQt5 import QtCore
from PyQt5.Qt import *
import sys
import os
import codecs
import copy

from libs.Dialog import LabelDialog, LabelDialog2
from libs.Scene import Scene
from libs.Generate_Xml import CreateXml, ReadXml


class Window(QMainWindow):
    def __init__(self, *args, **kwargs):
        super().__init__()

        self.setWindowTitle("数据标定")
        self.setWindowIcon(QIcon("icon/title.png"))
        self.setGeometry(100, 100, 600, 800)
        self.setMouseTracking(True)
        # 公共变量
        self.multiple = 1.0
        self.files_url = []
        self.num = 0
        self.scene = None
        self.label_clist = []
        self.list_item = None  # 停靠栏列表标签信息
        self.img_size = None  # 图片尺寸深度信息
        self.pic_dir = None  # 图片完整路径

        # 调用方法
        self.left_toolbar(False)
        self.top_menu()
        self.show_area()
        self.load_predefined_classes("D:\\Program Project\\PythonProject\\labeltool\\labeltoolv0.7\\data\\tagsdata.txt")
        self.dock()
        self.info_ini()

        # signal link slot

    def closeEvent(self, *args, **kwargs):
        '''重写关闭窗口事件'''
        self.save_ini()
        return

    def save_ini(self):
        '''生成配置信息文件'''
        settings = QSettings("config", QSettings.IniFormat)  # 实例化QSettings控件并传入相应参数
        settings.setValue("window_size", self.size())  # 存储主窗口尺寸
        settings.setValue("window_coordinates", self.pos())  # 存储主窗口坐标

    def info_ini(self):
        '''使用配置文件'''
        settings = QSettings("config", QSettings.IniFormat)
        if settings.value("window_size") is None:
            self.setGeometry(400, 200, 800, 800)
            print("config文件还未生成，请关闭一次窗口以生成config文件")
        else:
            self.move(settings.value("window_coordinates"))
            self.resize(settings.value("window_size"))

    def show_dialog(self):
        '''创建对话框,以及对话框中信号链接slot'''
        dialog = LabelDialog(self, "输入标签：", self.label_clist)
        dialog.add_rec.connect(self.scene.create_rect)
        dialog.add_rec.connect(self.add_label)
        dialog.disadd_rec.connect(self.scene.disadd_rec)

    def load_predefined_classes(self, tagsdata):
        '''从预分类标签文件中读取标签信息'''
        if os.path.exists(tagsdata) is True:
            with codecs.open(tagsdata, 'r') as f:
                for line in f:
                    line = line.strip()
                    if self.label_clist is None:
                        self.label_clist = [line]
                    else:
                        self.label_clist.append(line)

    def left_toolbar(self, bool):
        '''左侧工具栏'''
        self.toolbar = QToolBar()  # 创建工具栏
        self.addToolBar(Qt.LeftToolBarArea,
                        self.toolbar)  # addToolBar使用QMainWindow类的方法添加一个新的工具栏，参数为（位置，工具栏对象）,因此要先创建一个QToolBar控件对象，再使用addToolBar方法添加这个工具栏对象到指定位置，如果不指定位置则会默认在窗口上方位置，addToolBar(self, str) -> QToolBar
        self.toolbar.setToolButtonStyle(Qt.ToolButtonTextUnderIcon)  # 设置工具栏风格为按钮文本在图标下方
        self.toolbar.setIconSize(QtCore.QSize(40, 40))  # 调整图标大小

        # 创建行为动作
        open_button = QAction(QIcon("icon/文件.png"), "打开文件", self)
        opdr_button = QAction(QIcon("icon/文件.png"), "打开目录", self)
        transdr_button = QAction(QIcon("icon/文件.png"), "改变存放目录", self)
        preimage_button = QAction(QIcon("icon/前一个.png"), "上一张图像", self)
        nextimage_button = QAction(QIcon("icon/下一个.png"), "下一张图像", self)
        self.save_button = QAction(QIcon("icon/保存.png"), "保存", self)
        self.delete_button = QAction(QIcon("icon/删除.png"), "删除", self)
        self.amplification_button = QAction(QIcon("icon/放大.png"), "放大", self)
        self.lessen_button = QAction(QIcon("icon/缩小.png"), "缩小", self)
        self.origin_button = QAction(QIcon("icon/原始尺寸.png"), "原始尺寸", self)  # 原始尺寸返回的是图片原大小，而不是打开图片时的显示的大小
        self.pain_button = QAction(QIcon("icon/标注.png"), "标注", self)
        self.hadjus_button = QAction(QIcon("icon/高度自适应.png"), "适应窗口高度", self)
        self.wadjus_button = QAction(QIcon("icon/宽度自适应.png"), "适应窗口宽度", self)

        # 链接slot函数
        open_button.triggered.connect(self.op_file)
        opdr_button.triggered.connect(self.op_dir)
        preimage_button.triggered.connect(self.pre_image)
        nextimage_button.triggered.connect(self.next_image)
        self.save_button.triggered.connect(self.save_data)
        self.delete_button.triggered.connect(self.delete_item)
        self.amplification_button.triggered.connect(lambda: self.scale_img(1.25))  # 链接slot函数时无法传参，所以链接匿名函数，再用匿名函数调用
        self.lessen_button.triggered.connect(lambda: self.scale_img(0.8))
        self.origin_button.triggered.connect(self.origin_size)
        self.pain_button.triggered.connect(self.painter)
        self.hadjus_button.triggered.connect(self.height_adjust)
        self.wadjus_button.triggered.connect(self.width_adjust)

        # 添加行为到工具栏
        self.toolbar.addAction(open_button)
        self.toolbar.addAction(opdr_button)
        self.toolbar.addAction(transdr_button)
        self.toolbar.addAction(preimage_button)
        self.toolbar.addAction(nextimage_button)
        self.toolbar.addAction(self.save_button)
        self.toolbar.addAction(self.delete_button)
        self.toolbar.addSeparator()
        self.toolbar.addAction(self.amplification_button)
        self.toolbar.addAction(self.lessen_button)
        self.toolbar.addAction(self.origin_button)
        self.toolbar.addAction(self.pain_button)
        self.toolbar.addAction(self.hadjus_button)
        self.toolbar.addAction(self.wadjus_button)

        self.action_bool(bool)

    def action_bool(self, bool):
        '''工具栏按钮使能'''
        self.save_button.setEnabled(bool)
        self.delete_button.setEnabled(bool)
        self.amplification_button.setEnabled(bool)
        self.lessen_button.setEnabled(bool)
        self.origin_button.setEnabled(bool)
        self.pain_button.setEnabled(bool)
        self.hadjus_button.setEnabled(bool)
        self.wadjus_button.setEnabled(bool)

    def dock(self):
        '''右侧停靠栏'''
        self.dock = QDockWidget("区块的标签")  # 创建停靠栏
        self.file_dock = QDockWidget("文件列表")
        self.addDockWidget(Qt.RightDockWidgetArea, self.dock)  # 将停靠栏添加到主窗口指定区域
        self.addDockWidget(Qt.RightDockWidgetArea, self.file_dock)
        self.file_dock.setFeatures(QDockWidget.DockWidgetFloatable)
        self.dock.setFeatures(QDockWidget.DockWidgetMovable)
        # self.dock.setMinimumSize(200,200)
        # self.file_dock.setMinimumSize(200,100)

        # 设置区块标签停靠栏
        # 水平布局的两个控件
        self.use_default_label_checkbox = QCheckBox('使用预设标签')  # 复选框
        self.use_default_label_checkbox.setChecked(False)  # 设置为未选状态
        self.default_label_combo_box = QComboBox(self)  # 下拉框

        use_default_label_qhbox_layout = QHBoxLayout()  # 水平布局
        use_default_label_qhbox_layout.addWidget(self.use_default_label_checkbox)
        use_default_label_qhbox_layout.addWidget(self.default_label_combo_box)
        use_default_label_container = QWidget()  # 创建容器，容纳上面的复选框和下拉栏
        use_default_label_container.setLayout(use_default_label_qhbox_layout)  # 给容器设置布局方式

        # 其余控件
        self.diffc_button = QCheckBox('有难度的')  # 复选框
        self.diffc_button.setChecked(False)
        self.combo_box = QComboBox(self)  # 下拉框
        self.label_list = QListWidget()  # 列表控件
        self.label_list.itemClicked.connect(self.label_select)
        self.label_list.itemDoubleClicked.connect(self.edit_label)
        label_list_container = QWidget()  # 容纳控件的容器

        list_layout = QVBoxLayout()
        list_layout.setContentsMargins(0, 0, 0, 0)

        list_layout.addWidget(self.diffc_button)
        list_layout.addWidget(use_default_label_container)
        list_layout.addWidget(self.combo_box)
        list_layout.addWidget(self.label_list)

        label_list_container.setLayout(list_layout)

        self.dock.setWidget(label_list_container)  ##setWidget()用于设置在DockWidget内显示的QWidget及其子类对象

        # 设置self.file_dock文件列表停靠栏
        self.file_list_widget = QListWidget()  # 创建列表控件
        file_list_layout = QVBoxLayout()  # 创建布局方式
        file_list_layout.setContentsMargins(5, 0, 0, 0)  # 设置布局方式边距
        list_widget_container = QWidget()  # 创建QWidget控件作为容器
        file_list_layout.addWidget(self.file_list_widget)
        list_widget_container.setLayout(file_list_layout)
        self.file_dock.setWidget(list_widget_container)  # setWidget()用于设置在DockWidget内显示的QWidget及其子类对象

    def add_label(self):
        '''将标签添加到停靠栏的列表控件中'''
        self.label_list.clear()
        for item, label in self.scene.Rect_item.items():
            QListWidgetItem(label[0], self.label_list, label[1])

    def remove_label(self):
        '''将标签从停靠栏列表中移除'''
        if self.list_item:  # 在未点击图元时，self.list_item为被赋值，选中图元后，self.list_item被赋值，对应于选中的图元，这里解决的是如何删除与图元对应的标签的问题
            row = self.label_list.row(self.list_item)
            self.label_list.takeItem(row)
        else:
            row = self.label_list.currentRow()
            self.label_list.takeItem(row)

    def label_select(self, list_item):
        '''设置选中标签的同时选中对应的图元'''
        print(list_item.type())
        for item, label in self.scene.Rect_item.items():
            item.setSelected(False)
            if label[1] == list_item.type():
                item.setSelected(True)

    def rect_select(self, new_item, old_item, reason):
        '''选中图元时选中标签'''
        item = new_item
        count = self.label_list.count()
        for i in range(count):
            if item in self.scene.Rect_item:
                if self.label_list.item(i).type() == self.scene.Rect_item[item][1]:
                    self.label_list.item(i).setSelected(True)
                    self.list_item = self.label_list.item(i)

    def edit_label(self, item):
        '''编辑标签'''
        dialog = LabelDialog2(self, "输入标签：", self.label_clist)

        def edit(str):
            item.setText(str)
            new_label = str
            for it, label in self.scene.Rect_item.items():
                if label[1] == item.type():
                    self.scene.Rect_item[it][0] = new_label

        dialog.saved.connect(edit)

    def delete_item(self):
        '''删除图元'''
        items = self.scene.selectedItems()
        if items is not None:
            for item in items:
                self.scene.removeItem(item)
                label = self.scene.Rect_item[item]
                self.remove_label()
                del self.scene.Rect_item[item]
                del item

        else:
            pass

    def painter(self):
        '''启动标注,slot'''
        self.scene.flage()

    # 图片自适应窗口高度,slot
    def height_adjust(self):
        pass

    # 图片自适应窗口宽度，slot
    def width_adjust(self):
        pass

    def origin_size(self):
        '''原始尺寸函数,slot'''
        self.view.scale(1 / self.multiple, 1 / self.multiple)
        self.multiple = 1.0

    def save_data(self):
        '''保存矩形坐标数据 slot'''
        if self.scene is not None:
            path = self.pic_dir
            filepath, filename = os.path.split(path)
            folder = os.path.basename(filepath)
            img_size = self.img_size
            label_list = self.scene.output()
            xml = CreateXml(folder, filename, path, img_size, label_list)
            xml.save()

    # 顶部菜单栏
    def top_menu(self):
        # 在主窗口实例化QMenuBar()菜单栏控件
        bar = self.menuBar()
        # 添加菜单
        menu1 = bar.addMenu("文件&F")
        menu2 = bar.addMenu("编辑&A")
        menu3 = bar.addMenu("查看&M")
        menu4 = bar.addMenu("帮助&H")

        # 行为动作创建
        open_file = QAction("打开文件", menu1)
        open_file.triggered.connect(self.op_file)

        open_directory = QAction("打开目录", menu1)
        trasnsfer_directory = QAction("改变存放目录", menu1)
        open_label = QAction("开启标签", menu1)
        copy_frame = QAction("复制当前图像中的上一个边界框", menu1)
        save = QAction("保存", menu1)
        Pascal_VOC = QAction("Pascal/VOC", menu1)
        transfer_save = QAction("另存为", menu1)
        close_file = QAction("关闭文件", menu1)
        reset = QAction("全部重置", menu1)
        delete_image = QAction("删除图像", menu1)
        exit = QAction("退出", menu1)

        create_block = QAction("创建区块", menu2)
        edit_label = QAction("编辑标签", menu2)
        copy_block = QAction("复制区块", menu2)
        delete_block = QAction("删除所选区块", menu2)
        block_lineclolor = QAction("区块线条颜色", menu2)
        paint_squarer = QAction("绘制正方形", menu2)

        auto_save = QAction("自动保存模式", menu3)
        singel_class_mode = QAction("单一类别模式", menu3)
        show_class = QAction("显示类别", menu3)
        show_hide_label = QAction("显示/隐藏标签", menu3)
        professional_mode = QAction("专家模式", menu3)
        hide_block = QAction("隐藏所有区块", menu3)
        show_block = QAction("显示所有区块", menu3)
        amplify_image = QAction("放大画面", menu3)
        amplify_image.triggered.connect(lambda: self.scale_img(1.25))
        lessen_image = QAction("缩小画面", menu3)
        lessen_image.triggered.connect(lambda: self.scale_img(0.8))
        origin_image = QAction("原始大小", menu3)
        origin_image.triggered.connect(self.origin_size)

        none = QAction("暂无", menu4)
        # 创建子菜单
        sub_menu1 = QMenu("最近打开", menu1)

        # 子菜单包含的行为动作
        sub_action = QAction("待定", sub_menu1)
        # 添加动作到子菜单
        sub_menu1.addAction(sub_action)

        # 添加动作到父菜单
        menu1.addAction(open_file)
        menu1.addAction(open_directory)
        menu1.addAction(trasnsfer_directory)
        menu1.addAction(open_label)
        menu1.addAction(copy_frame)
        menu1.addMenu(sub_menu1)
        menu1.addAction(save)
        menu1.addAction(Pascal_VOC)
        menu1.addAction(transfer_save)
        menu1.addAction(close_file)
        menu1.addAction(reset)
        menu1.addAction(delete_image)
        menu1.addAction(exit)

        menu2.addAction(create_block)
        menu2.addAction(edit_label)
        menu2.addAction(copy_block)
        menu2.addAction(delete_block)
        menu2.addSeparator()
        menu2.addAction(block_lineclolor)
        menu2.addAction(paint_squarer)

        menu3.addAction(auto_save)
        menu3.addAction(singel_class_mode)
        menu3.addAction(show_class)
        menu3.addAction(show_hide_label)
        menu3.addAction(professional_mode)
        menu3.addSeparator()
        menu3.addAction(hide_block)
        menu3.addAction(show_block)
        menu3.addSeparator()
        menu3.addAction(amplify_image)
        menu3.addAction(lessen_image)
        menu3.addAction(origin_image)

        menu4.addAction(none)

    def op_dir(self):
        '''打开目录时获取整个目录的文件路径'''
        self.num = 0
        self.dir = QFileDialog.getExistingDirectory(self, "打开图片目录", "./")
        self.files_url.clear()
        for root, dirs, files in os.walk(self.dir):
            for file in files:
                self.files_url.append(root + '/' + file)
        if len(self.files_url):
            self.show_image(self.files_url[self.num])
            self.pic_dir = self.files_url[self.num]

    def pre_image(self):
        '''上一张图片'''
        length = len(self.files_url)
        if length:
            if 0 < self.num < length:
                self.num -= 1
                self.show_image(self.files_url[self.num])
                self.pic_dir = self.files_url[self.num]

    def next_image(self):
        '''下一张图片'''
        length = len(self.files_url)
        if length:
            if 0 <= self.num < length - 1:
                self.num += 1
                self.show_image(self.files_url[self.num])
                self.pic_dir = self.files_url[self.num]

    # 创建文件对话框,获取文件路径,slot
    def op_file(self):
        '''获取文件路径'''
        self.filename = QFileDialog.getOpenFileName(self, "打开图片", "./", "图片(*.png *.jpg)")
        # 对文件对话框内容做判断，判断是否打开了文件
        if self.filename[0] == "":  # 判断是否为空字符串
            pass
        else:
            self.show_image(self.filename[0])  # 调用打开图片方法
            self.pic_dir = self.filename[0]

    def show_area(self):
        '''view'''
        self.view = QGraphicsView()
        self.setCentralWidget(self.view)

    def show_image(self, url):
        '''显示图片，并且调用按钮使能函数'''
        image = QImage(url)
        self.img_size = [image.width(), image.height(), 1 if image.isGrayscale() else 3]
        pixmap = QPixmap.fromImage(image)
        self.scene = Scene(self.view.width(), self.view.height(), pixmap)
        self.view.setScene(self.scene)
        self.action_bool(True)
        self.label_from_xml(url)

        self.scene.complete_rect.connect(self.show_dialog)
        self.scene.focusItemChanged.connect(self.rect_select)

    def scale_img(self, factor):
        '''图片缩放函数,slot'''
        self.multiple *= factor  # 记录缩放的倍数
        self.view.scale(factor, factor)

    def label_from_xml(self, file_path):
        '''从xml导入标注'''
        xml = ReadXml(file_path)
        picture_size = xml.size()
        item_information = xml.information()
        if (picture_size is not None) and (item_information is not None):
            self.scene.rect_generate_from_xml(picture_size, item_information)
            self.add_label()


if __name__ == '__main__':
    app = QApplication(sys.argv)

    window = Window()
    window.show()

    sys.exit(app.exec_())
